package com.zhj.leetcode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.PriorityQueue;

/**
 * 力扣 292 前K个高频单词
 */
public class Test692 {
    public static void main(String[] args) {
        String[] words = {"i", "love", "leetcode", "i", "love", "coding"};
        System.out.println(toKFrequent(words, 2));
    }

    private static List<String> toKFrequent(String[] words, int k) {
        HashMap<String, Integer> map = new HashMap<>();
        for (String word : words) {
            if (!map.containsKey(word)) {
                map.put(word, 1);
            } else {
                map.put(word, map.get(word) + 1);
            }
        }
        PriorityQueue<String> pq  = new PriorityQueue<>(
                (w1, w2) -> {
                    return map.get(w1).equals(map.get(w2)) ? w1.compareTo(w2) : map.get(w1) - map.get(w2);
                }
        );

        for (String word : map.keySet()) {
            pq.add(word);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        List<String> res = new ArrayList<>();
        while (!pq.isEmpty()) {
            res.add(pq.poll());
        }
        return res;
    }
}
